"""
利用angr进行数据处理，获得functionSim的模型输入
"""
import os
from tqdm import tqdm
import angr
import networkx as nx
import shelve
import datetime
import os
from tqdm import tqdm
import time
import signal


def set_timeout(num, callback):
    """
    函数超时设置，避免一直卡在一个函数中
    """
    def wrap(func):
        # 收到信号 SIGALRM 后的回调函数，第一个参数是信号的数字，第二个参数是the interrupted stack frame.
        def handle(signum, frame):
            raise RuntimeError

        def to_do(*args, **kwargs):
            try:
                signal.signal(signal.SIGALRM, handle)  # 设置信号和回调函数
                signal.alarm(num)  # 设置 num 秒的闹钟
                print('start alarm signal.')
                r = func(*args, **kwargs)
                print('close alarm signal.')
                signal.alarm(0)  # 关闭闹钟
                return r
            except RuntimeError as e:
                callback()
        return to_do
    return wrap


def after_timeout():  # 超时后的处理函数
    print("Time out!")
    return False, 0


timeOut = 600


def getUnfinishedName(originPath, savePath):
    sampleList = {}
    finishList = {}
    unfinList = {}
    failedList = {}
    sLoc = 0
    fLoc = 0
    uLoc = 0
    # 获取所有的样本名
    for file in os.listdir(originPath):
        data = file.split(".")  # 找到样本文件
        if (len(data) == 1):
            sLoc += 1
            sampleList[data[0]] = True
    # 获取已经处理过的样本名
    for file in os.listdir(savePath):
        data = file.split(".")
        if (data[1] == "bak"):
            fLoc += 1
            finishList[data[0]] = True
    # 获得不能处理的样本
    path = "/home/cyw/projects/function_sim_project/disassemblyTools/angr/failedSample.txt"
    with open(path) as file:
        for str in file.readlines():
            failedList[str.strip()] = True
    # 获得未被处理过的样本
    for name in sampleList.keys():
        if name not in finishList and name not in failedList:
            uLoc += 1
            unfinList[name] = True

    logs("样本总计{}个，已处理{}个,还剩{}个\n".format(sLoc, fLoc, uLoc))
    return unfinList


logPath = r"/home/cyw/projects/function_sim_project/disassemblyTools/angr/angr_logs.txt"


def logs(inf):
    with open(logPath, "a+") as file:
        file.write(str(datetime.datetime.now())+"\t" + inf + "\n")


class angrDevTools:
    dataTrans = ["mov", "push", "pop", "xchg", "in", "out", "xlat",
                 "cwd", "lea", "lds", "les", "lahf", "sahf", "pushf", "popf"]
    arithmetric = ["add", "adc", "inc", "daa", "sub", "sbb", "dec", "cmp",
                   "das", "mul", "imul", "div", "idiv", "xor", "or", "not", "sll", "srl"]
    controlTrans = ["call", "ret", "jmp", "loop", "loope", "loopne", "int", "into", "iret", "jz",
                    "je", "jnz", "jne", "jo", "jno", "jc", "jnc", "js", "jns", "jb", "jnae", "jnb", "jae", "ja",
                    "jnbe", "jna", "jbe", "jl", "jnge", "jnl", "jge", "jg", "jnle", "jng", "jle"]

    theadNum = 1
    flag = False
    res = {}

    def __init__(self) -> None:
        fcgPath = r'E:\code\share\functionSImData\predata\angr_ACG'
        savePath = r'E:\code\share\functionSImData\predata\angr_tensor_acg'
        pass

    def calc_function(self, function):
        """
            提取函数的特征
            0：字符常量数量
            1：数值常量数量
            2：转移指令数量
            3：控制指令数量
            4：指令数量
            5：算术指令数量
            6：后代数量
            7：上一辈数量
        """
        funcData = dict()
        funcData["name"] = function.name  # 函数名
        funcData["features"] = [0, 0, 0, 0, 0, 0, 0, 0]  # 函数的特征属性
        funcData["type"] = "local"  # 是否为静态链接库函数
        # logs("生成函数："+str(funcData["name"]))
        try:
            res = [0, 0, 0, 0, 0, 0, 0, 0]
            # 函数中字符串常量的个数
            # 去看源码，返回是一个生成器（yield (md.addr,md.content)），不能直接使用len函数
            for tar in function.string_references():
                # logs("字符串：\t"+str(tar))
                res[0] += 1

            #   the set of basic blocks belonging to the function, that you can explore and disassemble using capstone.
            #   提取函数内每个基本块的属性并统计基本块的个数
            for blk in function.blocks:
                # 数字常量的个数
                res[1] = len(blk.vex.constants)
                transNum, controlNum, arithmeticNum = 0, 0, 0
                num_calls = 0
                for ins in blk.capstone.insns:
                    ins_name = ins.insn_name()
                    if ins_name in self.dataTrans:
                        transNum = transNum + 1
                    if ins_name in self.arithmetric:
                        arithmeticNum = arithmeticNum + 1
                    if ins_name in self.controlTrans:
                        controlNum = controlNum + 1
            res[2] = transNum
            res[3] = controlNum
            res[5] = arithmeticNum
            res[4] = blk.instructions
            res[6] = 0
            res[7] = 0
            funcData["features"] = res
        except Exception as e:
            logs(str(function.name)+"  Exception->"+str(e))
            funcData["type"] = "static link"
        return funcData

    def save_data(self, res, savePath):
        file = shelve.open(savePath)
        file["cg"] = res["cg"]
        file["cgattr"] = res["cgattr"]
        file["functype"] = res["functype"]
        file["funcs_id"] = res["funcs_id"]
        file.close()

    def use_angr_parse(self, sample):
        logs("正在处理样本\t"+str(sample))
        try:
            p = angr.Project(sample, load_options={'auto_load_libs': False})
            cfg = p.analyses.CFGFast()
            callGraph = cfg.functions.callgraph
        except Exception as e:
            logs("Exception->"+str(e))
            logs(str(sample)+"汇编错误")
            return False, 0, 0
        return True, cfg, callGraph

    # 分析文件夹下的样本，
    def analysisPath(self, originPath, savePath):
        unfinishedList = getUnfinishedName(originPath, savePath)
        logs("使用angr,进行样本反汇编")
        for name in tqdm(unfinishedList):
            samplePath = originPath + "/" + name
            self.analysisBinary(samplePath)
            if (self.flag):
                self.save_data(self.res, savePath + "/" + name)
            else:
                path = "/home/cyw/projects/function_sim_project/disassemblyTools/angr/failedSample.txt"
                with open(path, "a+") as file:
                    file.write(name+"\n")
                logs("angr处理失败")
        logs("样本反汇编完成")

    @set_timeout(timeOut, after_timeout)
    def analysisBinary(self, sample):
        """
            进行样本的分析
            保存如下数据：
            # cg
            # cgattr
            # functype
            # funcs_id
        """
        self.flag = False
        self.res = {}
        flag, cfg, callGraph = self.use_angr_parse(sample)
        if flag == False:
            self.flag = False
            self.res = {}
            return
        # 获得邻接矩阵
        cg = nx.adjacency_matrix(callGraph).todense().tolist()
        # 对应函数名和邻接矩阵的下标
        funcs_id = {}
        functype = {}
        cgattr = []
        res = {}
        num = 0
        nodeNum = len(cg)
        for func in cfg.kb.functions.values():
            funcTar = cfg.kb.functions[func.addr]
            funcData = self.calc_function(funcTar)
            # is_plt显然不对，is_syscall可能是,也不是
            # find_declaration也不是，查关键字statically linked binaries无结果。。。
            # 函数类型可能有问题
            # logs(str(funcTar.find_declaration())+"---"+str(funcTar.calling_convention)+"---"+ \
            #       str(funcTar.prototype_libname)+"---"+ str(funcTar.is_plt))
            if (funcTar.find_declaration()):
                funcData["type"] = "dynamic import"
            feature = funcData["features"]
            for i in range(nodeNum):
                if (cg[num][i] != 0):
                    feature[6] += 1
            for j in range(nodeNum):
                if cg[j][num] != 0:
                    feature[7] += 1
            functype[num] = funcData["type"]
            funcs_id[funcData["name"]] = num
            cgattr.append(feature)
            num += 1
        self.res["cg"] = cg
        self.res["funcs_id"] = funcs_id
        self.res["cgattr"] = cgattr
        self.res["functype"] = functype
        self.flag = True


if __name__ == '__main__':
    originPath = r"/home/cyw/projects/data/functionsim/allorigindata"
    savePath = r"/home/cyw/projects/function_sim_project/all_data/sampleDatas/angr"
    angrTools = angrDevTools()
    angrTools.analysisPath(originPath, savePath)
